
import torch

# LambdaLR 用于自定义 学习率 调整
from torch.optim.lr_scheduler import LambdaLR

import learn_conf as lc

# 主要是验证模型的正确性
# 模型的任务是复制输入，训练目标和输入是相同的序列。
# 拷贝任务是一个简单的序列到序列（seq2seq）学习任务，模型的目标是把输入的符号序列完整无误地输出。
# 测试模型的基础记忆和序列建模能力很有用。


# 生成 copy-task 的数据
# @param (V) : 词汇表大小，代表序列中符号的种类数（编号从 1 到 V-1)
# @param (batch_size) : 一个 batch 中有多少的 seq
# @param (nbatches) : 生成多少个 batch
def data_gen(V, batch_size, nbatches):
    "Generate random data for a src-tgt copy task."
    # 生成每个 batch 中的数据
    for i in range(nbatches):
        # 从 [1, V) 中 随机取数字，生成一个 (batch_size, seq_len) 的 tensor
        # batch_size 就是 一个 batch 中有多少 seq, 10 就是这一个 seq 中有多少 token
        data = torch.randint(1, V, size=(batch_size, 10))
        # 将所有序列的第一个元素（第 0 列）都设置成 1。 1 是一个 起始标签
        data[:, 0] = 1
        # 这里必须使用 clone(), 如果不使用 clone() 的话, 会导致 src 和 tgt 指向同一块内存
        src = data.requires_grad_(False).clone().detach()
        tgt = data.requires_grad_(False).clone().detach()
        # yield 是用来逐步的生成 数据
        # 会返回一个 迭代器 对象，当这个对象使用 next() 方法的时候才会继续运行获取下一个值
        yield Batch(src, tgt, 0)


def train_copy_task():
    V = 11
    criterion = LabelSmoothing(size=V, padding_idx=0, smoothing=0.0)
    model = make_model(V, V, N=2)

    # 创建 optimizer 并且 自定义 学习率
    optimizer = torch.optim.Adam(
        model.parameters(), lr=0.5, betas=(0.9, 0.98), eps=1e-9
    )
    lr_scheduler = LambdaLR(
        optimizer=optimizer,
        lr_lambda=lambda step: lc.learningRate(
            step, model_size=model.src_embed[0].d_model, factor=1.0, warmup=400
        ),
    )

    batch_size = 80
    for epoch in range(20):
        model.train()
        run_epoch(
            data_gen(V, batch_size, 20),
            model,
            SimpleLossCompute(model.generator, criterion),
            optimizer,
            lr_scheduler,
            mode="train",
        )
        model.eval()
        run_epoch(
            data_gen(V, batch_size, 5),
            model,
            SimpleLossCompute(model.generator, criterion),
            DummyOptimizer(),
            DummyScheduler(),
            mode="eval",
        )[0]

    model.eval()
    src = torch.LongTensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]])
    max_len = src.shape[1]
    src_mask = torch.ones(1, 1, max_len)
    print(greedy_decode(model, src, src_mask, max_len=max_len, start_symbol=0))

